/*
 * This file is part of yapk (yet another (auto) process killer) driver.
 *
 * Copyright (C) 2011 by Mateusz Wojcik <mewojcik@gmail.com>
 *
 * This is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation; either version 2.1 of
 * the License, or (at your option) any later version.
 *
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this software; if not, write to the Free
 * Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
 * 02110-1301 USA, or see the FSF site: http://www.fsf.org.
 */
#include <ntddk.h>
#include <stdlib.h>
#include <stdio.h>
#include "yapk_drv.h"
#include "utils.h"
#include "../../common/contract.h"


PDEVICE_OBJECT		DeviceObjectPtr;					// Ptr to newly created device object.

/**
 * \brief This is driver entry routine.
 * @param [in] DriverObjectPtr ptr to driver object created by Io Manager.
 * @param [in] RegistryPathPtr ptr to registry path string created for this driver.
 * @return NT_STATUS_SUCCESS when everything goes ok, !NT_STATUS_SUCCESS otherwise.
 */
NTSTATUS DriverEntry(	__in PDRIVER_OBJECT		DriverObjectPtr,
						__in PUNICODE_STRING	RegistryPathPtr )
{
	NTSTATUS			OperationStatus;					// Holds operations status.
	UNICODE_STRING		NtDeviceName;						// Holds internal (not-exported) driver name.
	UNICODE_STRING		DosDeviceName;						// Holds external (exported) driver name.	
	PYapkDevExtension	DeviceExtensionPtr;					// Ptr to driver extension.

	// [1] Supress warning - no need to use registry path.
	UNREFERENCED_PARAMETER(RegistryPathPtr);

	// [2] Initialize strings.
	RtlInitUnicodeString(&NtDeviceName,		NT_DEVICE_NAME	);
	RtlInitUnicodeString(&DosDeviceName,	DOS_DEVICE_NAME	);

	// [3] Try to create device object for this driver.
	OperationStatus = IoCreateDevice(	DriverObjectPtr,
										SIZEOF_YAPK_DEV_EXTENSION,
										&NtDeviceName,
										FILE_DEVICE_UNKNOWN,
										FILE_DEVICE_SECURE_OPEN,
										FALSE,
										&DeviceObjectPtr);

	if (STATUS_SUCCESS != OperationStatus)
	{
		DRV_ERROR("IoCreateDevice", OperationStatus);
		return OperationStatus;
	}

	// [4] Prepare dispatch table for drv.
	DriverObjectPtr->DriverUnload							= DriverUnload;
	DriverObjectPtr->MajorFunction[IRP_MJ_CREATE]			= DriverCreateClose;
	DriverObjectPtr->MajorFunction[IRP_MJ_CLOSE]			= DriverCreateClose;
	DriverObjectPtr->MajorFunction[IRP_MJ_DEVICE_CONTROL]	= DriverDeviceControl;

	// [5] Create symlink for user space app.
	OperationStatus = IoCreateSymbolicLink(&DosDeviceName, &NtDeviceName);
	if (STATUS_SUCCESS != OperationStatus)
	{
		DRV_ERROR("IoCreateSymbolicLink", OperationStatus);
		IoDeleteDevice(DeviceObjectPtr);
		return OperationStatus;
	}

	// [6] Initialize extension structures.
	DeviceExtensionPtr = (PYapkDevExtension)DeviceObjectPtr->DeviceExtension;
	InitializeListHead(&(DeviceExtensionPtr->ProcessNameListHead));
	ExInitializeFastMutex(&(DeviceExtensionPtr->ProcessNameListGuard));

	// [7] Register callback for process creation/termination event.
	OperationStatus = PsSetCreateProcessNotifyRoutine(PsNotifyCreateProcessRoutine, FALSE);
	if (STATUS_SUCCESS != OperationStatus)
	{
		DRV_ERROR("PsSetCreateProcessNotifyRoutine", OperationStatus);
		IoDeleteDevice(DeviceObjectPtr);
		IoDeleteSymbolicLink(&DosDeviceName);
		return OperationStatus;
	}

	DRV_DBG(("Hurrey! Becoming alive!\n"));

	return STATUS_SUCCESS;
}


/**
 * \brief Unloads driver.
 * @param [in] DriverObjectPtr - ptr to driver object.
 */
VOID DriverUnload(__in PDRIVER_OBJECT DriverObjectPtr)
{
	UNICODE_STRING	DosDeviceName;		// Holds external (exported) device name.
	
	// [1] Initialize string.
	RtlInitUnicodeString(&DosDeviceName, DOS_DEVICE_NAME);

	// [2] Remove callback for notify process creation event.
	PsSetCreateProcessNotifyRoutine(PsNotifyCreateProcessRoutine, TRUE /*remove*/);

	// [3] Remove symlink.
	IoDeleteSymbolicLink(&DosDeviceName);

	// [4] Remove device if exists.
	if (NULL != DriverObjectPtr->DeviceObject)
	{
		IoDeleteDevice(DriverObjectPtr->DeviceObject);
	}

	DRV_DBG(("Goodbye cruel world!\n"));
}

/**
 * \brief This function is responsible for simply completing 
 *        request when device is opened or closed.
 * @param DeviceObjectPtr ptr to device obejct
 * @param IrpPtr ptr to request.
 * @return STATUS_SUCCESS in case of success !STATUS_SUCCESS otherwise.
 */
NTSTATUS DriverCreateClose(PDEVICE_OBJECT DeviceObjectPtr, PIRP IrpPtr)
{
	// [1] Supress warnings
	UNREFERENCED_PARAMETER(DeviceObjectPtr);

	// [2] Complete request.
	DRV_COMPLETE_REQUEST(IrpPtr, STATUS_SUCCESS, 0, IO_NO_INCREMENT);

	return STATUS_SUCCESS;
}

/**
 * \brief This is main control function.
 * @param DeviceObjectPtr ptr to device object.
 * @param IrpPtr ptr to request.
 * @return STATUS_SUCCESS in case of success, !STATUS_SUCCESS otherwise.
 */
NTSTATUS DriverDeviceControl(PDEVICE_OBJECT DeviceObjectPtr, PIRP IrpPtr)
{
	PIO_STACK_LOCATION		IrpStackLocation;			// Holds ptr to io stack location.
	ULONG					InputBufferLength;			// Holds input buffer length.
	ULONG					OutputBufferLength;			// Holds output buffer length;
	PYapkDevExtension		DeviceExtensionsPtr;		// Holds device extensions pointer.
	PCHAR					InputBufferPtr;				// Holds input buffer.
	PCHAR					OutputBufferPtr;			// Holds output buffer.
	PLIST_ENTRY				EntryPtr;					// Holds entries in killer list.
	PYapkProcessEntry		ProcessEntryPtr;			// Holds process entry.
	PYapkProcessInfoList	ProcessListPtr;				// Holds process info list.

	// [1] Get all parameters from requestor.
	IrpStackLocation	= IoGetCurrentIrpStackLocation(IrpPtr);
	InputBufferLength	= IrpStackLocation->Parameters.DeviceIoControl.InputBufferLength;
	OutputBufferLength	= IrpStackLocation->Parameters.DeviceIoControl.OutputBufferLength;
	DeviceExtensionsPtr	= (PYapkDevExtension)DeviceObjectPtr->DeviceExtension;

	// [2] Check ioctl code.
	switch(IrpStackLocation->Parameters.DeviceIoControl.IoControlCode)
	{
		case IOCTL_ADD_PROCESS_TO_BE_KILLED:
			/*Adds process to killer list.*/

			// [3] Get buffers.
			InputBufferPtr	= (PCHAR)IrpPtr->AssociatedIrp.SystemBuffer;
			OutputBufferPtr	= (PCHAR)IrpPtr->AssociatedIrp.SystemBuffer;
			if (0 == InputBufferLength || NULL == InputBufferPtr || NULL == OutputBufferPtr)
			{
				DRV_ERROR("sanity check", STATUS_INSUFFICIENT_RESOURCES);
				DRV_COMPLETE_REQUEST(IrpPtr, STATUS_INSUFFICIENT_RESOURCES, 0, IO_NO_INCREMENT);
				return STATUS_INSUFFICIENT_RESOURCES;
			}

			// [4] Allocate list entry from non-paged pool.
			ProcessEntryPtr = (PYapkProcessEntry)ExAllocatePoolWithQuotaTag(PagedPool, SIZEOF_YAPK_PROCESS_ENTRY, ALLOCATION_TAG);

			// [5] Fill process entry structure.
			RtlCopyBytes(ProcessEntryPtr->ProcessName,InputBufferPtr, InputBufferLength); // change it!!! (InputBufferLength > STRING_MAX) ? STRING_MAX : InputBufferLength
			
			// [6] Insert entry into the list.
			ExAcquireFastMutex(&(DeviceExtensionsPtr->ProcessNameListGuard));
				InsertTailList(&(DeviceExtensionsPtr->ProcessNameListHead), &(ProcessEntryPtr->Entry));
			ExReleaseFastMutex(&(DeviceExtensionsPtr->ProcessNameListGuard));
		break;

		case IOCTL_DEL_PROCESS_TO_BE_KILLED:
			/*Removes process from killer list.*/

			// [3] Get buffers.
			InputBufferPtr	= (PCHAR)IrpPtr->AssociatedIrp.SystemBuffer;
			OutputBufferPtr	= (PCHAR)IrpPtr->AssociatedIrp.SystemBuffer;
			if (0 == InputBufferLength || NULL == InputBufferPtr || NULL == OutputBufferPtr)
			{
				DRV_ERROR("sanity check", STATUS_INSUFFICIENT_RESOURCES);
				DRV_COMPLETE_REQUEST(IrpPtr, STATUS_INSUFFICIENT_RESOURCES, 0, IO_NO_INCREMENT);
				return STATUS_INSUFFICIENT_RESOURCES;
			}

			// [4] Interate thru list and try to find matching process on list.
			ExAcquireFastMutex(&(DeviceExtensionsPtr->ProcessNameListGuard));
				
				EntryPtr = DeviceExtensionsPtr->ProcessNameListHead.Flink;
				while (EntryPtr != &(DeviceExtensionsPtr->ProcessNameListHead))
				{

					ProcessEntryPtr	= CONTAINING_RECORD(EntryPtr, YapkProcessEntry, Entry);

					if (0 == strcmp(ProcessEntryPtr->ProcessName, InputBufferPtr))
					{	/*got match!*/

						// [5] Remove element from list.
						RemoveEntryList(EntryPtr);
						ExFreePoolWithTag(ProcessEntryPtr, ALLOCATION_TAG);
						break;
					}

					EntryPtr = EntryPtr->Flink;
				} // while

			ExReleaseFastMutex(&(DeviceExtensionsPtr->ProcessNameListGuard));
		break;

		case IOCTL_SHOW_KILLER_LIST:
			/*Returns killer list to user space app.*/

			/*
			 * This is temporary solution, will be changed to non buffered method
			 * of passing input and output data between user-space and kernel-space.
			 */

			// [1] Get input and output buffers.
			ProcessListPtr	= (PYapkProcessInfoList)IrpPtr->AssociatedIrp.SystemBuffer;
			ProcessListPtr	= (PYapkProcessInfoList)IrpPtr->AssociatedIrp.SystemBuffer;
			if (0 == InputBufferLength || NULL == ProcessListPtr)
			{
				DRV_ERROR("sanity check", STATUS_INSUFFICIENT_RESOURCES);
				DRV_COMPLETE_REQUEST(IrpPtr, STATUS_INSUFFICIENT_RESOURCES, 0, IO_NO_INCREMENT);
				return STATUS_INSUFFICIENT_RESOURCES;
			}

			// [2] Zero input/output buffer.
			RtlZeroMemory(ProcessListPtr, SIZEOF_YAPK_PROCESS_INFO_LIST);

			// [3] Fill list.
			ExAcquireFastMutex(&(DeviceExtensionsPtr->ProcessNameListGuard));

				EntryPtr = DeviceExtensionsPtr->ProcessNameListHead.Flink;
				while (EntryPtr != &(DeviceExtensionsPtr->ProcessNameListHead))
				{
					ProcessEntryPtr = CONTAINING_RECORD(EntryPtr, YapkProcessEntry, Entry);

					// [4] Copy content to list
					strcpy((ProcessListPtr->List[ProcessListPtr->NumOfEntries]).ProcessName, ProcessEntryPtr->ProcessName);

					// [5] Increment number of entries.
					ProcessListPtr->NumOfEntries++;

					EntryPtr = EntryPtr->Flink;
				}

			ExReleaseFastMutex(&(DeviceExtensionsPtr->ProcessNameListGuard));

			// [6] Complete request by our own.
			DRV_COMPLETE_REQUEST(IrpPtr, STATUS_SUCCESS, SIZEOF_YAPK_PROCESS_INFO_LIST, IO_NO_INCREMENT);
			return STATUS_SUCCESS;

		break;

		default:
			DRV_DBG(("Unknown ioctl - ignoring.\n"));
	}

	DRV_COMPLETE_REQUEST(IrpPtr, STATUS_SUCCESS, 0, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

VOID PsNotifyCreateProcessRoutine (IN HANDLE  ParentId, IN HANDLE  ProcessId, IN BOOLEAN  Create)
{
	CHAR				CreatedProcessName[STRING_MAX];						// Holds newly created process name.
	CLIENT_ID			ClientId;											// Holds client unique identifier.
	PYapkDevExtension	DeviceExtensionPtr				= NULL;				// Holds device extension.
	PLIST_ENTRY			EntryPtr						= NULL;				// List iterator.
	PYapkProcessEntry	ProcessEntryPtr					= NULL;				// Holds process information entry.
	BOOLEAN				OperationStatus					= FALSE;			// Holds operation status.
	NTSTATUS			Status							= STATUS_SUCCESS;	// Holds operation status.
	HANDLE				CreatedProcessObjectHandle		= NULL;				// Holds created process object handle.

	OBJECT_ATTRIBUTES	ProcessObjectAttributes			=					// Holds Process Object attributes
	{ 
		sizeof(OBJECT_ATTRIBUTES),											// Structure length.
		0,																	// Root directory.
		NULL,																// Object name.
		0,																	// Attributes.
		NULL,																// Security Descriptor.
		NULL																// Security QoS.
	};

	// [1] Supress warning.
	UNREFERENCED_PARAMETER(ParentId);

	// [2] Check if this is creation callback.
	if (TRUE == Create)
	{
		// [3] Get device extension.
		DeviceExtensionPtr = (PYapkDevExtension)DeviceObjectPtr->DeviceExtension;

		// [4] Get newly created process name.
		RtlZeroMemory(CreatedProcessName, STRING_MAX);
		OperationStatus = RetrieveProcessNameByID(ProcessId, CreatedProcessName, STRING_MAX);
		if (FALSE == OperationStatus)
		{
			DRV_ERROR("RetrieveProcessNameByID", -1);
			return;
		}

		// [5] Try to find process name on driver internal list.
		ExAcquireFastMutex(&(DeviceExtensionPtr->ProcessNameListGuard));

			EntryPtr = DeviceExtensionPtr->ProcessNameListHead.Flink;
			while (EntryPtr != &(DeviceExtensionPtr->ProcessNameListHead))
			{

				ProcessEntryPtr	= CONTAINING_RECORD(EntryPtr, YapkProcessEntry, Entry);

				if (0 == strcmp(ProcessEntryPtr->ProcessName, CreatedProcessName))
				{	/*Process found on killers list - going to kill it.*/

					// [6] Initialize client identifier.
					ClientId.UniqueProcess = ProcessId;
					ClientId.UniqueThread  = 0;

					// [7] Get Process Object Handle.
					Status = ZwOpenProcess(&CreatedProcessObjectHandle, GENERIC_ALL, &ProcessObjectAttributes, &ClientId);
					if (STATUS_SUCCESS != Status)
					{
						DRV_ERROR("ZwOpenProcess", Status);
						return;
					}

					// [8] Terminate process.
					Status = ZwTerminateProcess(CreatedProcessObjectHandle, STATUS_SUCCESS);
					if (STATUS_SUCCESS != Status)
					{
						DRV_ERROR("ZwTerminateProcess", Status);
						return;
					}

					// [9] Close opened handle.
					Status = ZwClose(CreatedProcessObjectHandle);
					if (STATUS_SUCCESS != Status)
					{
						DRV_ERROR("ZwClose", Status);
						return;
					}
				}

				EntryPtr = EntryPtr->Flink;
			}

		ExReleaseFastMutex(&(DeviceExtensionPtr->ProcessNameListGuard));
	}
}


